/*
 * @(#)$Id$$
 *
 * Copyright 2006-2008 Makoto YUI
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * 
 * Contributors:
 *     Google Inc.- initial implementation
 *     Makoto YUI - ported
 */
/*
 * Copyright (C) 2007 Google Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package gridool.util.concurrent.reference;

import static gridool.util.concurrent.reference.ReferenceType.STRONG;

import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.lang.ref.Reference;
import java.util.AbstractMap;
import java.util.AbstractSet;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;


/**
 * Concurrent hash map that wraps keys and/or values in soft or weak references.
 * Does not support null keys or values. Uses identity equality for weak and
 * soft keys.
 *
 * <p>The concurrent semantics of {@link ConcurrentHashMap} combined with the
 * fact that the garbage collector can asynchronously reclaim and clean up
 * keys and values at any time can lead to some racy semantics. For example,
 * {@link #size()} returns an upper bound on the size; that is, the actual size
 * may be smaller in cases where the key or value has been reclaimed but the map
 * entry has not been cleaned up yet.
 *
 * @author crazybob@google.com (Bob Lee)
 * @author fry@google.com (Charles Fry)
 * @author Makoto YUI (yuin405@gmail.com)
 */
@SuppressWarnings("unchecked")
public class ReferenceMap<K, V> extends AbstractMap<K, V>
        implements ConcurrentMap<K, V>, Serializable {

    transient ConcurrentMap<Object, Object> delegate;

    private final ReferenceType keyReferenceType;
    private final ReferenceType valueReferenceType;

    /**
     * Concurrent hash map that wraps keys and/or values based on specified
     * reference types.
     *
     * @param delegateMap internally used map instance
     * @param keyReferenceType key reference type
     * @param valueReferenceType value reference type
     */
    public ReferenceMap(ConcurrentMap delegateMap, ReferenceType keyReferenceType, ReferenceType valueReferenceType) {
        ensureNotNull(keyReferenceType, valueReferenceType);

        if(keyReferenceType == ReferenceType.PHANTOM || valueReferenceType == ReferenceType.PHANTOM) {
            throw new IllegalArgumentException("Phantom references not supported.");
        }

        this.delegate = delegateMap;
        this.keyReferenceType = keyReferenceType;
        this.valueReferenceType = valueReferenceType;
    }

    /**
     * Concurrent hash map that wraps keys and/or values based on specified
     * reference types.
     *
     * @param keyReferenceType key reference type
     * @param valueReferenceType value reference type
     */
    public ReferenceMap(ReferenceType keyReferenceType, ReferenceType valueReferenceType) {
        this(new ConcurrentHashMap(), keyReferenceType, valueReferenceType);
    }
    
    public ReferenceMap(ReferenceType keyReferenceType, ReferenceType valueReferenceType, int initSize) {
        this(new ConcurrentHashMap(initSize), keyReferenceType, valueReferenceType);
    }

    V internalGet(K key) {
        Object valueReference = delegate.get(makeKeyReferenceAware(key));
        return dereferenceValue(valueReference);
    }

    public V get(final Object key) {
        ensureNotNull(key);
        return internalGet((K) key);
    }

    private V execute(Strategy strategy, K key, V value) {
        ensureNotNull(key, value);
        Object keyReference = referenceKey(key);
        return (V) strategy.execute(this, keyReference, referenceValue(keyReference, value));
    }

    public V put(K key, V value) {
        return execute(putStrategy(), key, value);
    }

    public V remove(Object key) {
        ensureNotNull(key);
        Object referenceAwareKey = makeKeyReferenceAware(key);
        Object valueReference = delegate.remove(referenceAwareKey);
        return dereferenceValue(valueReference);
    }

    public int size() {
        return delegate.size();
    }

    public boolean isEmpty() {
        return delegate.isEmpty();
    }

    public boolean containsKey(Object key) {
        ensureNotNull(key);
        Object referenceAwareKey = makeKeyReferenceAware(key);
        return delegate.containsKey(referenceAwareKey);
    }

    public boolean containsValue(Object value) {
        ensureNotNull(value);
        for(Object valueReference : delegate.values()) {
            if(value.equals(dereferenceValue(valueReference))) {
                return true;
            }
        }
        return false;
    }

    public void putAll(Map<? extends K, ? extends V> t) {
        for(Map.Entry<? extends K, ? extends V> entry : t.entrySet()) {
            put(entry.getKey(), entry.getValue());
        }
    }

    public void clear() {
        delegate.clear();
    }

    public V putIfAbsent(K key, V value) {
        return execute(putIfAbsentStrategy(), key, value);
    }

    public boolean remove(Object key, Object value) {
        ensureNotNull(key, value);
        Object referenceAwareKey = makeKeyReferenceAware(key);
        Object referenceAwareValue = makeValueReferenceAware(value);
        return delegate.remove(referenceAwareKey, referenceAwareValue);
    }

    public boolean replace(K key, V oldValue, V newValue) {
        ensureNotNull(key, oldValue, newValue);
        Object keyReference = referenceKey(key);

        Object referenceAwareOldValue = makeValueReferenceAware(oldValue);
        return delegate.replace(keyReference, referenceAwareOldValue, referenceValue(keyReference, newValue));
    }

    public V replace(K key, V value) {
        return execute(replaceStrategy(), key, value);
    }

    private transient volatile Set<Map.Entry<K, V>> entrySet = null;

    public Set<Map.Entry<K, V>> entrySet() {
        if(entrySet == null) {
            entrySet = new EntrySet();
        }
        return entrySet;
    }

    /** Dereferences an entry. Returns null if the key or value has been gc'ed. */
    private Entry dereferenceEntry(Map.Entry<Object, Object> entry) {
        K key = dereferenceKey(entry.getKey());
        V value = dereferenceValue(entry.getValue());
        return (key == null || value == null) ? null : new Entry(key, value);
    }

    /** Creates a reference for a key. */
    Object referenceKey(K key) {
        switch(keyReferenceType) {
            case STRONG:
                return key;
            case SOFT:
                return new SoftKeyReference(key);
            case WEAK:
                return new WeakKeyReference(key);
            default:
                throw new AssertionError();
        }
    }

    /** Converts a reference to a key. */
    K dereferenceKey(Object o) {
        return (K) dereference(keyReferenceType, o);
    }

    /** Converts a reference to a value. */
    V dereferenceValue(Object o) {
        if(o == null) {
            return null;
        }
        Object value = dereference(valueReferenceType, o);
        if(o instanceof InternalReference) {
            InternalReference ref = (InternalReference) o;
            if(value == null) {
                // old value was garbage collected
                ref.finalizeReferent();
            }
        }
        return (V) value;
    }

    /** Returns the refererent for reference given its reference type. */
    private Object dereference(ReferenceType referenceType, Object reference) {
        return referenceType == STRONG ? reference : ((Reference) reference).get();
    }

    /** Creates a reference for a value. */
    Object referenceValue(Object keyReference, V value) {
        switch(valueReferenceType) {
            case STRONG:
                return value;
            case SOFT:
                return new SoftValueReference(keyReference, value);
            case WEAK:
                return new WeakValueReference(keyReference, value);
            default:
                throw new AssertionError();
        }
    }

    /**
     * Wraps key so it can be compared to a referenced key for equality.
     */
    private Object makeKeyReferenceAware(Object o) {
        return keyReferenceType == STRONG ? o : new KeyReferenceAwareWrapper(o);
    }

    /** Wraps value so it can be compared to a referenced value for equality. */
    private Object makeValueReferenceAware(Object o) {
        return valueReferenceType == STRONG ? o : new ReferenceAwareWrapper(o);
    }

    /**
     * Marker interface to differentiate external and internal references. Also
     * duplicates FinalizableReference and Reference.get for internal use.
     */
    interface InternalReference {
        void finalizeReferent();

        Object get();
    }

    private static int keyHashCode(Object key) {
        return System.identityHashCode(key);
    }

    /*
     * Tests weak and soft references for identity equality. Compares references
     * to other references and wrappers. If o is a reference, this returns true if
     * r == o or if r and o reference the same non-null object. If o is a wrapper,
     * this returns true if r's referent is identical to the wrapped object.
     */
    protected static boolean referenceEquals(Reference r, Object o) {
        // compare reference to reference.
        if(o instanceof InternalReference) {
            // are they the same reference? used in cleanup.
            if(o == r) {
                return true;
            }

            // do they reference identical values? used in conditional puts.
            Object referent = ((Reference) o).get();
            return referent != null && referent == r.get();
        }

        // is the wrapped object identical to the referent? used in lookups.
        return ((ReferenceAwareWrapper) o).unwrap() == r.get();
    }

    /**
     * Returns {@code true} if the specified value reference has been garbage
     * collected. The value behind the reference is also passed in, rather than
     * queried inside this method, to ensure that the return statement of this
     * method will still hold true after it has returned (that is, a value
     * reference exists outside of this method which will prevent that value from
     * being garbage collected).
     *
     * @param valueReference the value reference to be tested
     * @param value the object referenced by {@code valueReference}
     * @return {@code true} if {@code valueReference} is non-null and {@code
     *     value} is null
     */
    private static boolean isExpired(Object valueReference, Object value) {
        return (valueReference != null) && (value == null);
    }

    /**
     * Big hack. Used to compare keys and values to referenced keys and values
     * without creating more references.
     */
    static class ReferenceAwareWrapper {
        final Object wrapped;

        ReferenceAwareWrapper(Object wrapped) {
            this.wrapped = wrapped;
        }

        Object unwrap() {
            return wrapped;
        }

        public int hashCode() {
            return wrapped.hashCode();
        }

        public boolean equals(Object obj) {
            // defer to reference's equals() logic.
            return obj.equals(this);
        }
    }

    /** Used for keys. Overrides hash code to use identity hash code. */
    static class KeyReferenceAwareWrapper extends ReferenceAwareWrapper {
        public KeyReferenceAwareWrapper(Object wrapped) {
            super(wrapped);
        }

        public int hashCode() {
            return System.identityHashCode(wrapped);
        }
    }

    class SoftKeyReference extends FinalizableSoftReference<Object> implements InternalReference {
        final int hashCode;

        public SoftKeyReference(Object key) {
            super(key);
            this.hashCode = keyHashCode(key);
        }

        public void finalizeReferent() {
            delegate.remove(this);
        }

        @Override
        public int hashCode() {
            return this.hashCode;
        }

        @Override
        public boolean equals(Object o) {
            return referenceEquals(this, o);
        }
    }

    class WeakKeyReference extends FinalizableWeakReference<Object> implements InternalReference {
        final int hashCode;

        public WeakKeyReference(Object key) {
            super(key);
            this.hashCode = keyHashCode(key);
        }

        public void finalizeReferent() {
            delegate.remove(this);
        }

        @Override
        public int hashCode() {
            return this.hashCode;
        }

        @Override
        public boolean equals(Object o) {
            return referenceEquals(this, o);
        }
    }

    class SoftValueReference extends FinalizableSoftReference<Object> implements InternalReference {
        final Object keyReference;

        public SoftValueReference(Object keyReference, Object value) {
            super(value);
            this.keyReference = keyReference;
        }

        public void finalizeReferent() {
            delegate.remove(keyReference, this);
        }

        @Override
        public boolean equals(Object obj) {
            return referenceEquals(this, obj);
        }
    }

    class WeakValueReference extends FinalizableWeakReference<Object> implements InternalReference {
        final Object keyReference;

        public WeakValueReference(Object keyReference, Object value) {
            super(value);
            this.keyReference = keyReference;
        }

        public void finalizeReferent() {
            delegate.remove(keyReference, this);
        }

        @Override
        public boolean equals(Object obj) {
            return referenceEquals(this, obj);
        }
    }

    protected interface Strategy {
        public Object execute(ReferenceMap map, Object keyReference, Object valueReference);
    }

    // TODO(crazybob): a getter called put() is probably a bad idea
    protected Strategy putStrategy() {
        return PutStrategy.PUT;
    }

    protected Strategy putIfAbsentStrategy() {
        return PutStrategy.PUT_IF_ABSENT;
    }

    protected Strategy replaceStrategy() {
        return PutStrategy.REPLACE;
    }

    private enum PutStrategy implements Strategy {
        PUT {
            public Object execute(ReferenceMap map, Object keyReference, Object valueReference) {
                return map.dereferenceValue(map.delegate.put(keyReference, valueReference));
            }
        },

        REPLACE {
            public Object execute(ReferenceMap map, Object keyReference, Object valueReference) {
                // ensure that the existing value is not collected
                do {
                    Object existingValueReference;
                    Object existingValue;
                    do {
                        existingValueReference = map.delegate.get(keyReference);
                        existingValue = map.dereferenceValue(existingValueReference);
                    } while(isExpired(existingValueReference, existingValue));

                    if(existingValueReference == null) {
                        // nothing to replace
                        return false;
                    }

                    if(map.delegate.replace(keyReference, existingValueReference, valueReference)) {
                        // existingValue didn't expire since we still have a reference to it
                        return existingValue;
                    }
                } while(true);
            }
        },

        PUT_IF_ABSENT {
            public Object execute(ReferenceMap map, Object keyReference, Object valueReference) {
                Object existingValueReference;
                Object existingValue;
                do {
                    existingValueReference = map.delegate.putIfAbsent(keyReference, valueReference);
                    existingValue = map.dereferenceValue(existingValueReference);
                } while(isExpired(existingValueReference, existingValue));
                return existingValue;
            }
        },
    }

    private static PutStrategy defaultPutStrategy;

    protected PutStrategy getPutStrategy() {
        return defaultPutStrategy;
    }

    class Entry implements Map.Entry<K, V> {
        final K key;
        V value;

        public Entry(K key, V value) {
            this.key = key;
            this.value = value;
        }

        public K getKey() {
            return this.key;
        }

        public V getValue() {
            return this.value;
        }

        public V setValue(V newValue) {
            value = newValue;
            return put(key, newValue);
        }

        public int hashCode() {
            return key.hashCode() * 31 + value.hashCode();
        }

        public boolean equals(Object o) {
            if(!(o instanceof ReferenceMap.Entry)) {
                return false;
            }

            Entry entry = (Entry) o;
            return key.equals(entry.key) && value.equals(entry.value);
        }

        public String toString() {
            return key + "=" + value;
        }
    }

    private class ReferenceIterator implements Iterator<Map.Entry<K, V>> {
        private Iterator<Map.Entry<Object, Object>> i = delegate.entrySet().iterator();
        private Map.Entry<K, V> nextEntry;
        private Map.Entry<K, V> lastReturned;

        public ReferenceIterator() {
            advance();
        }

        private void advance() {
            while(i.hasNext()) {
                Map.Entry<K, V> entry = dereferenceEntry(i.next());
                if(entry != null) {
                    nextEntry = entry;
                    return;
                }
            }

            // nothing left
            nextEntry = null;
        }

        public boolean hasNext() {
            return nextEntry != null;
        }

        public Map.Entry<K, V> next() {
            if(nextEntry == null) {
                throw new NoSuchElementException();
            }
            lastReturned = nextEntry;
            advance();
            return lastReturned;
        }

        public void remove() {
            ReferenceMap.this.remove(lastReturned.getKey());
        }
    }

    private class EntrySet extends AbstractSet<Map.Entry<K, V>> {
        public Iterator<Map.Entry<K, V>> iterator() {
            return new ReferenceIterator();
        }

        public int size() {
            return delegate.size();
        }

        public boolean contains(Object o) {
            if(!(o instanceof Map.Entry)) {
                return false;
            }
            Map.Entry<K, V> e = (Map.Entry<K, V>) o;
            V v = ReferenceMap.this.get(e.getKey());
            return v != null && v.equals(e.getValue());
        }

        public boolean remove(Object o) {
            if(!(o instanceof Map.Entry)) {
                return false;
            }
            Map.Entry<K, V> e = (Map.Entry<K, V>) o;
            return ReferenceMap.this.remove(e.getKey(), e.getValue());
        }

        public void clear() {
            delegate.clear();
        }
    };

    // TODO(kevinb): use preconditions
    static void ensureNotNull(Object o) {
        if(o == null) {
            throw new NullPointerException();
        }
    }

    static void ensureNotNull(Object... array) {
        for(int i = 0; i < array.length; i++) {
            if(array[i] == null) {
                throw new NullPointerException("Argument #" + i + " is null.");
            }
        }
    }

    private void writeObject(ObjectOutputStream out) throws IOException {
        out.defaultWriteObject();
        out.writeInt(size());
        for(Map.Entry<Object, Object> entry : delegate.entrySet()) {
            Object key = dereferenceKey(entry.getKey());
            Object value = dereferenceValue(entry.getValue());

            // don't persist gc'ed entries.
            if(key != null && value != null) {
                out.writeObject(key);
                out.writeObject(value);
            }
        }
        out.writeObject(null);
    }

    private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
        in.defaultReadObject();
        int size = in.readInt();
        this.delegate = new ConcurrentHashMap<Object, Object>(size);
        while(true) {
            K key = (K) in.readObject();
            if(key == null) {
                break;
            }
            V value = (V) in.readObject();
            put(key, value);
        }
    }

    private static final long serialVersionUID = 0;
}
